import os
import cv2
import matplotlib.pyplot as plt
from skimage import data, draw, color, transform, feature, measure
import SimpleITK as sitk
import numpy as np
from scipy import misc
import pandas


def circle_detect(image):
    # 加载图片，转换成灰度图并检测边缘
    image_gray = color.rgb2gray(image)
    edges = feature.canny(image_gray, sigma=2.0)
    print(edges.shape)
    fig2, (ax1, ax2) = plt.subplots(ncols=2, nrows=1, figsize=(8, 4))

    ax1.set_title('Original picture')
    ax1.imshow(image)

    ax2.set_title('Edge (white) and result (red)')
    ax2.imshow(edges)

    plt.show()

    # 执行椭圆变换
    result = transform.hough_ellipse(edges, accuracy=20, threshold=250, min_size=100, max_size=120)
    print(result)
    result.sort(order='accumulator')  # 根据累加器排序

    # 估计椭圆参数
    best = list(result[-1])  # 排完序后取最后一个
    yc, xc, a, b = [int(round(x)) for x in best[1:5]]
    orientation = best[5]

    # 在原图上画出椭圆
    cy, cx = draw.ellipse_perimeter(yc, xc, a, b, orientation)
    image[cy, cx] = (0, 0, 255)  # 在原图中用蓝色表示检测出的椭圆

    # 分别用白色表示canny边缘，用红色表示检测出的椭圆，进行对比
    edges = color.gray2rgb(edges)
    edges[cy, cx] = (250, 0, 0)

    ax2.set_title('Edge (white) and result (red)')
    ax2.imshow(edges)

    plt.show()


def calculate_center(image):
    x = []
    y = []
    for i in range(image.shape[0]):
        for j in range(image.shape[1]):
            pixel = image[i, j]
            # next_pixel = image[i,j+1]
            # print(pixel,next_pixel)
            if pixel == True:
                y.append(j)
                x.append(i)
    x = np.mean(x)
    y = np.mean(y)
    print(x, y)
    return x, y


#
# def calculate_node(image):
#     count = 0
#     x_max = 0
#     y_max = 0
#     for i in range(image.shape[0]):
#         for j in range(image.shape[1]):
#             pixel = image[i,j]
#             # next_pixel = image[i,j+1]
#             # print(pixel,next_pixel)
#             if pixel==True:
#                 c = x+y
#                 if c>count:
#                     count = c
#                     x_max = i
#                     y_max = j
#     print(x_max,y_max)
#     return x_max,y_max

# 判断是否被包围
def is_around(image, x, y):
    # 横着被包围
    list_x0 = image[int(x), :int(y)]
    list_x1 = image[int(x), int(y):]
    # 竖着被包围
    list_y0 = image[int(x):, int(y)]
    list_y1 = image[:int(x), int(y)]
    print(sum(list_x0), sum(list_x1), sum(list_y0), sum(list_y1))
    if sum(list_x0) and sum(list_x1) and sum(list_y0) and sum(list_y1):
        return True
    else:
        return False


# 转换cv image
def cv_image(image):
    new_image = image[:, :, np.newaxis]
    for i in range(image.shape[0]):
        for j in range(image.shape[1]):
            pixel = image[i, j]
            # print(pixel)
            if pixel == 0.0:
                new_image[i, j] = 0
            else:
                new_image[i, j] = 255
    return new_image


def resample(image, transform):
    """
    This function resamples (updates) an image using a specified transform
    :param image: The sitk image we are trying to transform
    :param transform: An sitk transform (ex. resizing, rotation, etc.
    :return: The transformed sitk image
    """
    reference_image = image
    interpolator = sitk.sitkLinear
    default_value = 0
    return sitk.Resample(image, reference_image, transform,
                         interpolator, default_value)


def get_center(img):
    """
    This function returns the physical center point of a 3d sitk image
    :param img: The sitk image we are trying to find the center of
    :return: The physical center point of the image
    """
    width, height, depth = img.GetSize()
    return img.TransformIndexToPhysicalPoint((int(np.ceil(width / 2)),
                                              int(np.ceil(height / 2)),
                                              int(np.ceil(depth / 2))))


def rotation3d(image, theta_x, theta_y, theta_z, center_point=None, show=False):
    """
    This function rotates an image across each of the x, y, z axes by theta_x, theta_y, and theta_z degrees
    respectively
    :param image: An sitk MRI image
    :param theta_x: The amount of degrees the user wants the image rotated around the x axis
    :param theta_y: The amount of degrees the user wants the image rotated around the y axis
    :param theta_z: The amount of degrees the user wants the image rotated around the z axis
    :param show: Boolean, whether or not the user wants to see the result of the rotation
    :return: The rotated image
    """
    theta_x = np.deg2rad(theta_x)
    theta_y = np.deg2rad(theta_y)
    theta_z = np.deg2rad(theta_z)
    euler_transform = sitk.Euler3DTransform(get_center(image), theta_x, theta_y, theta_z, (0, 0, 0))

    if center_point:
        image_center = tuple(center_point)
    else:
        image_center = get_center(image)

    # image_center = get_center(image)

    # image_center = ()
    print(image_center)
    euler_transform.SetCenter(image_center)
    euler_transform.SetRotation(theta_x, theta_y, theta_z)
    resampled_image = resample(image, euler_transform)
    if show:
        plt.imshow(sitk.GetArrayFromImage(resampled_image)[:, :, 256])
        plt.show()
    return resampled_image


# 得到心脏的层
def obtain_slices(image, dim):
    if dim == 0:
        z, x, y = image.shape

        for i in range(z):
            image_slice = image[i]
            num = sum(sum(image_slice))
            if num > 0:
                start = i
                break
        for i in range(z - 1, 0, -1):
            image_slice = image[i]
            num = sum(sum(image_slice))
            if num > 0:
                end = i
                break
    elif dim == 1:
        z, x, y = image.shape

        for i in range(y):
            image_slice = image[:, :, i]
            num = sum(sum(image_slice))
            if num > 0:
                start = i
                break
        for i in range(y - 1, 0, -1):
            image_slice = image[:, :, i]
            num = sum(sum(image_slice))
            if num > 0:
                end = i
                break
    start = start + (end - start) // 3
    end = end - (end - start) // 3
    return start, end


def turn_view(start, end, imag, dim):
    angle_x, angle_y = [], []  # 准备保存心脏中间1/3部分的所有角度，为了最终求平均
    center_pointx = 0
    center_pointy = 0
    for i in range(start, end):
        if dim == 0:
            image_org = imag[i, :, :]
        elif dim == 1:
            image_org = imag[:, :, i]
        image = cv_image(image_org)  # 由于原始图像的像素点为1或0，而opencv必须接受像素点为255或0的图像，所以转换一下像素点
        cv2.imwrite("new_img.png", image)  # 临时存放一下当前slice的图像，为了以opencv的格式读入
        image = cv2.imread('new_img.png')  # 读入当前slice
        # 低于127时设置成255，255是白色，代表亮度（能量）最高，包括三通道的图像中(255,255,255)是白色
        ret, thresh = cv2.threshold(cv2.cvtColor(image, cv2.COLOR_BGR2GRAY), 127, 255, 0)
        contours, hierarchy = cv2.findContours(thresh, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)  # 得到轮廓信息

        if contours != []:
            cnt = contours[-1]  # 得到轮廓信息
            rect = cv2.minAreaRect(cnt)  # 最小轮廓矩形
            box = cv2.boxPoints(rect)  # 转换为四个点
            # x0, y0, x1, y1 = box  # box四个点
            # box[3] = [y1[0]+50,y1[1]+50]
            box = np.int64(box)  # 转换数据类型
            # 以下注释掉的四行是为了显示在图像上可视化出最小矩形
            # imgnew = cv2.drawContours(image, [box], 0, (255, 0, 0), 2)
            # cv2.imshow('boundingRect', imgnew)
            # #
            # cv2.waitKey(0)
            # cv2.destroyAllWindows()
            x0, y0, x1, y1 = box  # box四个点
            w, h = rect[1]  # rect中包含起始点，长宽信息
            x, y = rect[0]
            if h != 0.0:
                scale = w / h
                if dim == 0:
                    min_scale, max_scale = 0, 0.9
                else:
                    min_scale, max_scale = 1.5, 2.0

                # if scale<0.9:
                if scale > min_scale and scale < max_scale:
                    print('scale is available: ', scale)
                    angle_x.append(abs(y0[1] - x0[1]))  # 得到x轴的长度差
                    angle_y.append(abs(y0[0] - x0[0]))  # 得到y轴的长度差
                    if i == zshape // 2:
                        center_pointx = (x0[1] + y0[1]) / 2
                        center_pointy = (x0[0] + y0[0]) / 2
    xx = np.mean(angle_x)  # 平均多层
    yy = np.mean(angle_y)
    if dim == 0:
        if center_pointx:
            center_point = (center_pointx, center_pointy, zshape // 2)
        else:
            center_point = None
    elif dim == 1:
        center_point = None
    # print(xx, yy, center_point)
    theta = np.rad2deg(np.arctan2(yy, xx))  # 得到夹角
    return theta, center_point


if __name__ == '__main__':
    root = 'D:/code/FFR-heart/after_post_process/'  # 数据文件夹
    images = os.listdir(root)
    for image in images:  # 遍历文件夹中的图像
        print(image)
        image_name = image.split('_')[1][:5]  # 图像的名称是'_'分割之后的数字部分
        image_path = os.path.join(root, image)  # 图像的绝对路径
        labelImage = sitk.ReadImage(image_path)  # 读入原始数据
        imag = sitk.GetArrayFromImage(labelImage)  # 创建数组
        imag_results = imag.copy()  # 因为后边要更改数据中部分层次的像素点（cv_image(image_org)），所以保存一个副本
        zshape, xshape, yshape = imag_results.shape  # 得到数据的大小 z,x,y
        # print(imag_results.shape)  #(320,512,512)

        start_z, end_z = obtain_slices(imag, dim=0)  # 得到z轴（横断面）上心脏所在的slice的开始和结束序号
        print(start_z, end_z)
        theta_z, center_point = turn_view(start_z, end_z, imag, dim=0)
        print(theta_z, center_point)

        start_x, end_x = obtain_slices(imag, dim=1)
        theta_x, _ = turn_view(start_x, end_x, imag, dim=1)

        # 转动原图
        origin_imaging = 'test_data/case_' + image_name + '/imaging.nii.gz'
        labelImage = sitk.ReadImage(origin_imaging)  # in_file是nii.gz文件的路径

        image_out = rotation3d(labelImage, theta_x, 0, theta_z, center_point, show=False)
        # 创建保存nii数据的文件夹
        data_dir = os.path.join('center2')
        if not os.path.isdir(data_dir):
            os.makedirs(data_dir)

        # 写入nii数据
        sitk.WriteImage(image_out, os.path.join(data_dir, 'imaging_' + image_name + '_shortaxis.nii.gz'))

        # 转分割
        image_out = sitk.GetImageFromArray(imag_results)
        image_out = rotation3d(image_out, theta_x, 0, theta_z, center_point, show=False)

        # 写入nii数据
        sitk.WriteImage(image_out, os.path.join(data_dir, 'segmentation_' + image_name + '_shortaxis.nii.gz'))
